#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 19-5-26 21:25
# @Author  : JSQ
# @File    : tcpserver.py
# @Software: PyCharm
# socket_server传输文件服务端,监听并接收上报的各层状态

import os
import json
import struct
import socketserver
import traceback
import data_input
import runtime_input
import software_input
import network_input
import os_input
import MySQLdb
import hmac
import subprocess

#认证客户端链接
def conn_auth(conn):
    print('\n开始验证客户端发起的新连接的合法性:')
    msg=os.urandom(32)
    conn.sendall(msg)
    respone = conn.recv(16)
    with open(os.getcwd() + '/validusers.txt', 'r') as f:
        for line in f.readlines():
            secret_key = str.encode(line.strip())
            h = hmac.new(secret_key, msg)
            digest = h.digest()
            if hmac.compare_digest(respone,digest):
                return True
    return False
    f.close()

def data_handler(conn,bufsize=1024):
    if not conn_auth(conn):
        print('该连接不合法,关闭!')
        conn.close()
        return
    print('连接合法,开始通信!')
    conn.sendall(b'1')

#接收文件并保存
def recv_file(head_dir, tcp_client):
    buffsize = 1024
    filename = head_dir['filename']
    filesize_b = head_dir['filesize_bytes']
    recv_len = 0
    recv_mesg = b''
    f = open(os.getcwd() + '/JsonFileReport/' + filename, 'wb')
    while recv_len < filesize_b:
        if filesize_b - recv_len > buffsize:
            recv_mesg = tcp_client.recv(buffsize)
            recv_len += len(recv_mesg)
            f.write(recv_mesg)
        else:
            recv_mesg = tcp_client.recv(filesize_b - recv_len)
            recv_len += len(recv_mesg)
            f.write(recv_mesg)
    f.close()
    print('文件接收完成-->', filename)

raddr_sets = set()

#从BaseRequestHandler继承，并重写handle方法
class MyServer(socketserver.BaseRequestHandler):

    def handle(self):
        connection = self.request
        data_handler(connection)
        print("Connected by", connection)
        raddr_sets.add(connection.getpeername()[0])
        while True:                                                    # 循环监听（读取）来自客户端的数据
            try:
                struct_len = connection.recv(4)                        # 接受报头的长度
                struct_info_len = struct.unpack('i', struct_len)[0]    # 解析得到报头信息的长度
                head_info = connection.recv(struct_info_len)           # 接受报头的内容
                head_dir = json.loads(head_info.decode('utf-8'))       # 将报头的内容反序列化
                recv_file(head_dir, connection)
            except:
                if len(raddr_sets) == 1:
                    print("connected ip:", raddr_sets)
                    print("只有一台主机上报，请等待另外两台...")
                    break
                if len(raddr_sets) == 2:
                    print("connected ip:", raddr_sets)
                    print("只有两台主机上报，请等待另外一台...")
                    break
                if len(raddr_sets) == 3:
                    print("connected ip:", raddr_sets)
                    print("三台主机全部上报，可以开始存储数据库！")
                    raddr_sets.clear()
                    print("\n开始写入数据库...")                               # 将各层上报信息写入数据库
                    data_input.data()
                    runtime_input.runtime()
                    software_input.software()
                    network_input.network()
                    os_input.os()
                    subprocess.getoutput("sudo find ./JsonFileCmd -name '*.json'|xargs rm -r")    # 清空决策命令文件
                    # subprocess.getoutput("sudo find ./AttackReport -name '*.json'|xargs rm -r")   # 清空上报的攻击文件
                    print("\n完成一次攻击面变换，结果已显示在系统界面！###############")
                    break

def receive():
    HOST_S3 = '172.16.0.4'
    print("服务端开始监听#############################")
    s = socketserver.ThreadingTCPServer((HOST_S3, 8082), MyServer)     # 多线程监听
    s.serve_forever()                                                  # 服务器一直开着

if __name__ == '__main__':
    receive()